import fitz  # PyMuPDF
import re
import json
import requests
import logging
from typing import Optional, Dict, Any
import pandas as pd

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class LocalChatGLMClient:
    """本地部署的ChatGLM客户端"""

    def __init__(self, base_url: str = "http://localhost:8000", timeout: int = 120):
        """
        初始化本地ChatGLM客户端

        Args:
            base_url: 本地模型服务的URL
            timeout: 请求超时时间
        """
        self.base_url = base_url.rstrip('/')
        self.timeout = timeout
        self.chat_url = f"{self.base_url}/v1/chat/completions"
        self.generate_url = f"{self.base_url}/v1/completions"

    def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.1) -> str:
        """
        向本地ChatGLM模型发送生成请求

        Args:
            prompt: 输入提示
            max_tokens: 最大生成token数
            temperature: 温度参数

        Returns:
            模型生成的文本
        """
        try:
            # 尝试使用chat/completions接口（OpenAI兼容格式）
            payload = {
                "model": "chatglm3",
                "messages": [
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                "max_tokens": max_tokens,
                "temperature": temperature,
                "stream": False
            }
            logger.info(f"请求chat: {payload}")

            response = requests.post(
                self.chat_url,
                json=payload,
                timeout=self.timeout
            )

            if response.status_code == 200:
                result = response.json()
                return result['choices'][0]['message']['content']
            else:
                logger.warning(f"Chat接口失败: {response.status_code}, 尝试completions接口")
                return self._fallback_generate(prompt, max_tokens, temperature)

        except Exception as e:
            logger.error(f"ChatGLM请求失败: {e}")
            return self._fallback_generate(prompt, max_tokens, temperature)

    def _fallback_generate(self, prompt: str, max_tokens: int, temperature: float) -> str:
        """备用生成方法"""
        try:
            payload = {
                "prompt": prompt,
                "max_tokens": max_tokens,
                "temperature": temperature,
                "stream": False
            }

            response = requests.post(
                self.generate_url,
                json=payload,
                timeout=self.timeout
            )

            if response.status_code == 200:
                result = response.json()
                return result['choices'][0]['text']
            else:
                raise Exception(f"生成失败: {response.status_code}")

        except Exception as e:
            logger.error(f"备用生成方法也失败: {e}")
            return ""


class BillOfLadingExtractor:
    """提单数据提取器"""

    def __init__(self, chatglm_client: Optional[LocalChatGLMClient] = None):
        self.chatglm_client = chatglm_client

    def extract_from_pdf(self, pdf_path: str, use_chatglm: bool = True) -> Dict[str, Any]:
        """
        从PDF文件中提取提单数据

        Args:
            pdf_path: PDF文件路径
            use_chatglm: 是否使用ChatGLM进行增强提取

        Returns:
            提取的数据字典
        """
        logger.info(f"开始处理PDF文件: {pdf_path}")

        # 第一步：使用传统方法提取
        # traditional_result = self._extract_traditional(pdf_path)
        #
        # # 检查传统方法是否成功提取所有字段
        # if self._is_complete_result(traditional_result):
        #     logger.info("传统方法成功提取所有数据")
        #     return traditional_result

        logger.warning("传统方法未能完全提取数据，尝试ChatGLM增强提取")

        # 第二步：如果传统方法不完整且允许使用ChatGLM，则使用大模型
        # if use_chatglm and self.chatglm_client:
        return self._extract_with_chatglm(pdf_path)

            # 合并结果，优先使用传统方法的结果，用ChatGLM结果补充缺失字段
            # merged_result = self._merge_results(traditional_result, chatglm_result)

            # if self._is_complete_result(merged_result):
            #     logger.info("ChatGLM增强提取成功")
            #     return merged_result
            # else:
            #     logger.warning("ChatGLM增强提取仍未能获取完整数据")
            #     return merged_result
        # else:
        #     logger.warning("未使用ChatGLM或ChatGLM客户端未配置")
        #     return traditional_result

    def _extract_traditional(self, pdf_path: str) -> Dict[str, Any]:
        """使用传统方法提取数据"""
        try:
            doc = fitz.open(pdf_path)
            full_text = ""
            for page in doc:
                full_text += page.get_text()
            doc.close()

            return self._parse_text_traditional(full_text)

        except Exception as e:
            logger.error(f"传统提取方法失败: {e}")
            return {"packages": None, "weight": None, "volume": None, "error": str(e)}

    def _parse_text_traditional(self, text: str) -> Dict[str, Any]:
        """使用传统方法解析文本"""
        packages = None
        weight = None
        volume = None

        # 查找关键区域
        particulars_section = self._extract_particulars_section(text)

        if particulars_section:
            # 使用多种模式匹配
            packages = self._extract_packages(particulars_section)
            weight = self._extract_weight(particulars_section)
            volume = self._extract_volume(particulars_section)

        return {
            "packages": packages,
            "weight": weight,
            "volume": volume,
            "method": "traditional"
        }

    def _extract_particulars_section(self, text: str) -> Optional[str]:
        """提取PARTICULARS FURNISHED BY SHIPPER部分"""
        start_marker = "PARTICULARS FURNISHED BY SHIPPER"
        end_markers = ["FREIGHT & CHARGES", "ON BOARD DATE:", "FREIGHT PREPAID"]

        start_index = text.find(start_marker)
        if start_index == -1:
            return None

        # 查找最近的结束标记
        end_index = len(text)
        for marker in end_markers:
            marker_index = text.find(marker, start_index)
            if marker_index != -1 and marker_index < end_index:
                end_index = marker_index

        return text[start_index:end_index]

    def _extract_packages(self, text: str) -> Optional[int]:
        """提取包裹数量"""
        patterns = [
            r'(\d+)\s*CARTONS?',
            r'(\d+)\s*CTNS?',
            r'Packages?[\s:]*(\d+)',
            r'PCS?[\s:]*(\d+)'
        ]

        for pattern in patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                return int(match.group(1))
        return None

    def _extract_weight(self, text: str) -> Optional[float]:
        """提取重量"""
        patterns = [
            r'(\d+(?:\.\d+)?)\s*KGS?',
            r'Gross\s*Weight[\s:]*(\d+(?:\.\d+)?)',
            r'Weight[\s:]*(\d+(?:\.\d+)?)'
        ]

        for pattern in patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                return float(match.group(1))
        return None

    def _extract_volume(self, text: str) -> Optional[float]:
        """提取体积"""
        patterns = [
            r'(\d+(?:\.\d+)?)\s*CBM',
            r'Measurement[\s:]*(\d+(?:\.\d+)?)',
            r'Volume[\s:]*(\d+(?:\.\d+)?)',
            r'CBM[\s:]*(\d+(?:\.\d+)?)'
        ]

        for pattern in patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                return float(match.group(1))
        return None

    def _extract_with_chatglm(self, pdf_path: str) -> Dict[str, Any]:
        """使用ChatGLM提取数据"""
        try:
            doc = fitz.open(pdf_path)
            full_text = ""
            for page in doc:
                full_text += page.get_text()
            doc.close()

            return self.extract_with_chatglm_by_content(full_text)

        except Exception as e:
            logger.error(f"ChatGLM提取失败: {e}")
            return {"packages": None, "weight": None, "volume": None, "method": "chatglm", "error": str(e)}

    def extract_with_chatglm_by_content(self, full_text):
        prompt = self._build_chatglm_prompt(full_text)
        response = self.chatglm_client.generate(prompt)
        print(response)
        return self._parse_chatglm_response(response)

    def _build_chatglm_prompt(self, text: str) -> str:
        """构建ChatGLM提示词"""

        # return f"提取提单信息: {text}"
        return f"""请从以下提单文本中精确提取信息：

{text}

请提取以下字段：
1. "MBL/HBL No." - 提单号，通常字段名为"EXPORT REFERENCES"、"Bill of Lading No." 等(字符串)
2. "Container No." - 集装箱号(字符串)
3. "Invoice No." - 出口日期（格式：YYYY/MM/DD）
4. "Date of Exportation" - 发货人公司名称(字符串)
5. "Shipper/Exporter Company" - 发货人地址(字符串)
6. "Consignee Company" - 收货人公司名称  (字符串)
7. "Consignee Address" - 收货人地址(字符串)
8. "packages" - 包裹数量（仅数字）
9. "weight" - 重量（仅数字，单位：公斤）
10. "volume" - 体积（仅数字，单位：立方米）


请返回纯JSON格式，不要添加任何解释文字："""
# 请返回纯JSON格式，不要添加任何解释文字：
# {{ "hblNo." : 提单号
# "containerNo" : 集装箱号
#  "dateOfExportation" : 出口日期（格式：YYYY/MM/DD）
# "exporter[Company]" : 发货人公司名称
# "exporter[Address]" : 发货人地址
# "exporter[zipCode]" : 邮政编码
# "consignee[Company]" : 收货人公司名称
# "consignee[Address]" : 收货人地址
# "packages" : 包裹数量（仅数字）
# "weight" : 重量（仅数字，单位：公斤）
#  "volume" : 体积（仅数字，单位：立方米）}}
#
# 如果找不到某个数值，请将其设置为null。"""

    def _parse_chatglm_response(self, response: str) -> Dict[str, Any]:
        """解析ChatGLM响应"""
        try:
            # 清理响应文本，提取JSON部分
            cleaned_response = response.strip()

            # 查找JSON对象
            json_match = re.search(r'\{[^{}]*"packages"[^{}]*"weight"[^{}]*"volume"[^{}]*\}', cleaned_response)
            if json_match:
                json_str = json_match.group()
                result = json.loads(json_str)
                result["method"] = "chatglm"
                return result

            # 如果找不到完整JSON，尝试分别提取数值
            packages = self._extract_value_from_response(response, "packages")
            weight = self._extract_value_from_response(response, "weight")
            volume = self._extract_value_from_response(response, "volume")

            return {
                "packages": packages,
                "weight": weight,
                "volume": volume,
                "method": "chatglm"
            }

        except Exception as e:
            logger.error(f"解析ChatGLM响应失败: {e}")
            return {"packages": None, "weight": None, "volume": None, "method": "chatglm", "error": str(e)}

    def _extract_value_from_response(self, response: str, field: str) -> Optional[float]:
        """从响应文本中提取特定字段的数值"""
        patterns = [
            f'"{field}"[\\s:]*([\\d.]+)',
            f"{field}[\\s:]*([\\d.]+)",
            f"([\\d.]+)[\\s]*(?:是|为)[\\s]*{field}"
        ]

        for pattern in patterns:
            match = re.search(pattern, response, re.IGNORECASE)
            if match:
                try:
                    return float(match.group(1))
                except ValueError:
                    continue
        return None

    def _is_complete_result(self, result: Dict[str, Any]) -> bool:
        """检查结果是否完整"""
        return (result.get("packages") is not None and
                result.get("weight") is not None and
                result.get("volume") is not None)

    def _merge_results(self, traditional: Dict[str, Any], chatglm: Dict[str, Any]) -> Dict[str, Any]:
        """合并传统方法和ChatGLM的结果"""
        merged = traditional.copy()
        merged["method"] = "merged"

        # 用ChatGLM结果补充缺失的字段
        for field in ["packages", "weight", "volume"]:
            if merged.get(field) is None and chatglm.get(field) is not None:
                merged[field] = chatglm[field]

        return merged

def update_excel(excel_path:str, data):
    """更新Excel文件"""
    try:
        # 读取Excel文件
        with pd.ExcelFile(excel_path) as xls:
            # 更新Commercial Invoice工作表
            df_invoice = pd.read_excel(xls, 'Commercial Invoice', header=None)

            # 更新字段
            # MBL/HBL No.
            df_invoice.iloc[1, 1] = data.get('mbl_hbl_no', '')
            # Container No.
            df_invoice.iloc[2, 1] = data.get('container_no', '')
            # Date of Exportation
            df_invoice.iloc[2, 6] = data.get('date_of_exportation', '')
            # Shipper Company
            df_invoice.iloc[4, 1] = data.get('shipper_company', '')
            # Shipper Address
            df_invoice.iloc[5, 1] = data.get('shipper_address', '')
            # Consignee Company
            df_invoice.iloc[4, 5] = data.get('consignee_company', '')
            # Consignee Address
            df_invoice.iloc[5, 5] = data.get('consignee_address', '')

            # 更新Packing List工作表
            df_packing = pd.read_excel(xls, 'Packing List', header=None)

            # 保存更新后的Excel
            with pd.ExcelWriter(excel_path, engine='openpyxl') as writer:
                df_invoice.to_excel(writer, sheet_name='Commercial Invoice', index=False, header=False)
                df_packing.to_excel(writer, sheet_name='Packing List', index=False, header=False)

            print("Excel文件更新成功！")

    except Exception as e:
        print(f"Excel更新错误: {e}")


# 使用示例
def extractPdf(content):
    # 初始化ChatGLM客户端（根据您的本地部署调整URL）
    chatglm_client = LocalChatGLMClient(
        base_url="http://localhost:8000",  # 修改为您的本地模型服务地址
        timeout=600
    )

    # 初始化提取器
    extractor = BillOfLadingExtractor(chatglm_client)

    # 提取数据
    # pdf_path = "D:\\存档-待核对\\RFCU4069145-完结\\BL-OOLU2162031000.pdf"
    # pdf_path = "D:\\存档-待核对\\GAOU6252103-已做\\BL-149506063292.pdf"
    result = extractor.extract_with_chatglm_by_content(content)

    # 输出结果
    print("提取结果:")
    print(json.dumps(result, indent=2, ensure_ascii=False))

    # 生成所需的JSON格式
    final_result = {
        "packages": result.get("packages"),
        "weight": result.get("weight"),
        "volume": result.get("volume")
    }

    print("\n最终JSON格式:")
    print(json.dumps(final_result, ensure_ascii=False))
    return final_result


# 使用示例
def extractPdfInfo(pdf_path):
    # 初始化ChatGLM客户端（根据您的本地部署调整URL）
    chatglm_client = LocalChatGLMClient(
        base_url="http://localhost:8000",  # 修改为您的本地模型服务地址
        timeout=60
    )

    # 初始化提取器
    extractor = BillOfLadingExtractor(chatglm_client)

    # 提取数据
    # pdf_path = "D:\\存档-待核对\\RFCU4069145-完结\\BL-OOLU2162031000.pdf"
    # pdf_path = "D:\\存档-待核对\\GAOU6252103-已做\\BL-149506063292.pdf"
    result = extractor.extract_from_pdf(pdf_path, use_chatglm=True)

    # 输出结果
    print("提取结果:")
    print(json.dumps(result, indent=2, ensure_ascii=False))

    # 生成所需的JSON格式
    final_result = {
        "packages": result.get("packages"),
        "weight": result.get("weight"),
        "volume": result.get("volume")
    }

    print("\n最终JSON格式:")
    print(json.dumps(final_result, ensure_ascii=False))
    return final_result


if __name__ == "__main__":
    pdf_path = "D:\\存档-待核对\\RFCU4069145-完结\\BL-OOLU2162031000.pdf"
    # pdf_path = "D:\\存档-待核对\\GAOU6252103-已做\\BL-149506063292.pdf"
    extractPdfInfo(pdf_path)